import torch
import torch.nn as nn
from transformers import BertModel


class MultiSourceStockModel(nn.Module):
    def __init__(self):
        super().__init__()

        # ----------------- 模块定义 -----------------
        # 1. 新闻事件编码器 (BERT)
        self.news_encoder = BertModel.from_pretrained('bert-base-chinese')
        self.news_proj = nn.Linear(768, 64)  # 降维

        # 2. 市场情绪时序网络 (BiLSTM)
        self.sentiment_lstm = nn.LSTM(
            input_size=1,  # 情绪得分单特征
            hidden_size=32,
            bidirectional=True,
            batch_first=True
        )

        # 3. 宏观经济编码器 (全连接)
        self.macro_encoder = nn.Sequential(
            nn.Linear(10, 32),  # 假设10个宏观指标
            nn.ReLU(),
            nn.Linear(32, 32)
        )

        # 4. 价格序列网络 (Transformer)
        self.price_transformer = nn.TransformerEncoder(
            encoder_layer=nn.TransformerEncoderLayer(d_model=5, nhead=5),  # OHLCV
            num_layers=3
        )
        self.price_proj = nn.Linear(5, 64)

        # 5. 基本面编码器 (结构化数据处理)
        self.fundamental_net = nn.Sequential(
            nn.Linear(20, 64),  # PE,PB,ROE等20个基本面指标
            nn.LayerNorm(64),
            nn.ReLU()
        )

        # ----------------- 融合与预测 -----------------
        # 跨模态注意力融合
        self.cross_attn = nn.MultiheadAttention(embed_dim=64, num_heads=4)

        # 联合预测头
        self.fc = nn.Sequential(
            nn.Linear(64 * 5, 256),
            nn.GELU(),
            nn.Dropout(0.2),
            nn.Linear(256, 3)  # 输出: 涨/跌/持平的3分类
        )

    def forward(self, news, sentiment, macro, price, fundamental):
        # ----------------- 各模态特征提取 -----------------
        # 1. 新闻事件 (batch_size, seq_len, 768) → (batch,64)
        news_feat = self.news_encoder(**news).last_hidden_state[:, 0, :]
        news_feat = self.news_proj(news_feat)

        # 2. 市场情绪时序 (batch, seq_len, 1) → (batch,64)
        sent_out, _ = self.sentiment_lstm(sentiment)
        sent_feat = sent_out.mean(dim=1)  # 时序平均

        # 3. 宏观经济 (batch, 10) → (batch,32)
        macro_feat = self.macro_encoder(macro)

        # 4. 价格序列 (batch, seq_len, 5) → (batch,64)
        price_feat = self.price_transformer(price)
        price_feat = self.price_proj(price_feat[:, -1, :])  # 取最后时间步

        # 5. 基本面 (batch, 20) → (batch,64)
        fund_feat = self.fundamental_net(fundamental)

        # ----------------- 跨模态融合 -----------------
        # 拼接所有特征 (batch, 64+32+64+32+64=256)
        combined = torch.cat([
            news_feat,
            sent_feat,
            macro_feat,
            price_feat,
            fund_feat
        ], dim=1)

        # 跨模态注意力交互
        attn_out, _ = self.cross_attn(
            combined.unsqueeze(1),
            combined.unsqueeze(1),
            combined.unsqueeze(1)
        attn_out = attn_out.squeeze(1)

        # ----------------- 预测输出 -----------------
        return self.fc(attn_out)